#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4

##
## Copyright (C)  2007 - George Y. Kussumoto (georgeyk.dev@gmail.com)
##
## This program is free software; you can redistribute it and/or
## modify it under the terms of the GNU Lesser General Public License
## as published by the Free Software Foundation; either version 2
## of the License, or (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU Lesser General Public License for more details.
##
## You should have received a copy of the GNU Lesser General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., or visit: http://www.gnu.org/.
##
##
"""Pack and unpack data inside a PNG file"""

import base64
from math import sqrt
import optparse
import os.path
import sys

try:
    from PIL import Image
except ImportError:
    raise SystemExit("PIL module not found!")

SEP = "###"

#
# Helper functions
#

def _add_file_name_in_data_string(input_file, data_string):
    # filename###data_string
    file_name = os.path.basename(input_file)
    file_name += SEP
    return file_name + data_string

def _get_string_data(data):
    string_data = ""
    try:
        string_data = "".join([chr(d) for d in data])
    except TypeError:
        string_data = ""

    return string_data

def _get_size(total):
    assert total > 0

    start = int(sqrt(total))
    for i in xrange(start, 1, -1):
        if total % i == 0:
            # (width, height)
            return (total/i, i)

def _check_file(file_name, parser):
    if not os.path.isfile(file_name):
        parser.error(u"%s: File not found!" % file_name)
    else:
        return True

#
# Main functions
#

def pack(input_file, output_file):
    data = open(input_file, "rb").read()
    if not data:
        raise SystemExit("%s: is empty" % input_file)

    data_string = base64.encodestring(data)
    data_string = _add_file_name_in_data_string(input_file, data_string)
    image = Image.new("L", _get_size(len(data_string)))
    image.fromstring(data_string)
    image.save(output_file)

def unpack(input_file):
    try:
        image = Image.open(input_file)
    except IOError:
        raise SystemExit("%s: not an image file" % input_file)

    string_data = _get_string_data(image.getdata())
    if not string_data:
        raise SystemExit("%s: does not contain valid data" % input_file)

    output_file, sep, original_data = string_data.partition(SEP)
    out = open(output_file, "wb")
    out.write(base64.decodestring(original_data))
    out.close()

def main(args):
    usage = "Usage: %prog [options] file"
    parser = optparse.OptionParser(usage=usage)
    parser.add_option('-p', '--pack', dest='packfile',
                      help='Packs a given file')
    parser.add_option('-o', '--output', dest='output_file',
                      help='A file name to save the packed file - Optional')
    parser.add_option('-u', '--unpack', dest='unpackfile',
                      help='Unpacks a given file')
    options, args = parser.parse_args(args)

    if options.packfile and not options.unpackfile:
        input_file = options.packfile
        if options.output_file:
            output_file = options.output_file
            if not output_file.endswith(".png"):
                output_file += ".png"
        else:
            output_file = os.path.basename(input_file) + ".png"

        if _check_file(input_file, parser):
            pack(input_file, output_file)
    elif options.unpackfile and not options.output_file:
        input_file = options.unpackfile
        if _check_file(input_file, parser):
            unpack(input_file)
    else:
        parser.error("Use --help parameter for some help")

    return 0

if __name__ == '__main__':
    sys.exit(main(sys.argv))
